#include "ObjectDetection.h"

#ifdef WIN32
#include <io.h>
#else
#include <dirent.h>
#endif

using namespace cv;
using namespace std;

void ReadModelDirectory(const string& modelDirectory, vector<string>& models, bool addDirectoryName)
{
    models.clear();

#ifdef WIN32
    struct _finddata_t s_file;
    string str = modelDirectory + "\\*.*";

    intptr_t h_file = _findfirst(str.c_str(), &s_file);
    if (h_file != static_cast<intptr_t>(-1.0))
    {
        do
        {
            if (addDirectoryName)
                models.push_back(modelDirectory + "\\" + s_file.name);
            else
                models.push_back((string)s_file.name);
        } while (_findnext( h_file, &s_file ) == 0);
    }
    _findclose(h_file);
#else
    DIR* dir = opendir(modelDirectory.c_str());
    if (dir != NULL)
    {
        struct dirent* dent;
        while ((dent = readdir(dir)) != NULL)
        {
            if (addDirectoryName)
                models.push_back(modelDirectory + "/" + string(dent->d_name));
            else
                models.push_back(string(dent->d_name));
        }
    }
#endif

    sort(models.begin(), models.end());
}

void DetectObjects(const Mat& image, const string& modelDirectory, set<string>& tags)
{
    vector<LatentSvmDetector::ObjectDetection> objectDetections;

    vector<string> models;
    ReadModelDirectory(modelDirectory, models);

    LatentSvmDetector detector(models);
    if (detector.empty())
    {
        cout << "Models cann't be loaded" << endl;
        exit(-1);
    }
    vector<string> classNames = detector.getClassNames();
    detector.detect(image, objectDetections);

    const float confidenceLevel = 0.5;
    for (vector<LatentSvmDetector::ObjectDetection>::size_type i = 0; i != objectDetections.size(); ++i)
    {
        LatentSvmDetector::ObjectDetection od = objectDetections[i];
        if (od.score > confidenceLevel || abs(od.score - confidenceLevel) < 1e-6)
        {
            tags.insert(classNames[od.classID]);
        }
    }
}

void DetectObjects(const Mat& image, const string& modelDirectory, vector<ObjectData>& datas)
{
    datas.clear();
    vector<LatentSvmDetector::ObjectDetection> objectDetections;

    vector<string> models;
    ReadModelDirectory(modelDirectory, models);

    LatentSvmDetector detector(models);
    if (detector.empty())
    {
        cout << "Models cann't be loaded" << endl;
        exit(-1);
    }
    vector<string> classNames = detector.getClassNames();
    detector.detect(image, objectDetections);

    const float confidenceLevel = 0.5;
    for (vector<LatentSvmDetector::ObjectDetection>::size_type i = 0; i != objectDetections.size(); ++i)
    {
        LatentSvmDetector::ObjectDetection od = objectDetections[i];
        if (od.score > confidenceLevel || abs(od.score - confidenceLevel) < 1e-6)
        {
            datas.push_back(ObjectData(classNames[od.classID], od.rect));
        }
    }
}